import numpy as np
import os

def edit_npz(file_path):
    """
    加载 NPZ 文件，访问 'circuits' 键下的数据，
    允许用户选择并删除其中的一个电路条目，然后保存更改。

    Args:
        file_path (str): NPZ 文件的路径。
    """
    if not os.path.exists(file_path):
        print(f"错误：文件 '{file_path}' 不存在。")
        return

    original_data_handle = None # 用于确保文件句柄总是被关闭
    try:
        # 加载 NPZ 文件
        original_data_handle = np.load(file_path, allow_pickle=True)
        
        # 检查 'circuits' 键是否存在
        if 'circuits' not in original_data_handle:
            print(f"错误：在文件 '{file_path}' 中找不到 'circuits' 键。")
            keys_available = list(original_data_handle.keys())
            if keys_available:
                 print(f"可用的键有: {', '.join(keys_available)}")
            else:
                 print("文件似乎是空的或者没有可识别的键。")
            return

        # 访问 'circuits' 键下的数据，假设它是一个包含电路数据的字典
        # 使用 .item() 来获取存储在0维数组中的对象
        try:
            circuits_data = original_data_handle['circuits'].item()
        except AttributeError:
             print(f"错误：无法从 'circuits' 键中提取数据项。它可能不是预期的格式。类型为: {type(original_data_handle['circuits'])}")
             return
             
        if not isinstance(circuits_data, dict):
            print(f"错误：'circuits' 键下的数据不是一个字典。实际类型: {type(circuits_data)}")
            return

        keys = list(circuits_data.keys())

        if not keys:
            print(f"在 'circuits' 键下没有找到任何电路数据。")
            return

        print(f"文件 '{file_path}' 的 'circuits' 键下包含以下电路：")
        for i, key in enumerate(keys):
             # 尝试打印每个电路数据的类型或形状，如果它是 numpy 数组
            circuit_item = circuits_data[key]
            info = f"类型: {type(circuit_item)}"
            if isinstance(circuit_item, np.ndarray):
                info = f"形状: {circuit_item.shape}"
            print(f"{i + 1}: {key} ({info})")


        while True:
            try:
                choice = input(f"请输入要删除的电路的编号 (1-{len(keys)})，或输入 'q' 退出：")
                if choice.lower() == 'q':
                    print("操作已取消。")
                    return
                
                choice_index = int(choice) - 1
                if 0 <= choice_index < len(keys):
                    key_to_delete = keys[choice_index]
                    break
                else:
                    print("无效的选择，请输入列表中的编号。")
            except ValueError:
                print("无效的输入，请输入数字编号。")

        confirm = input(f"确定要删除电路 '{key_to_delete}' 吗？ (y/n): ")
        if confirm.lower() == 'y':
            # 从内部字典中删除选定的电路
            del circuits_data[key_to_delete]
            
            # 创建一个新的顶层字典，包含修改后的 circuits_data
            # 如果原 NPZ 文件除了 'circuits' 外还有其他顶级键需要保留，需在此处添加
            output_data = {'circuits': circuits_data} 
            
            # 确保关闭旧的文件句柄后才写入
            original_data_handle.close()
            original_data_handle = None # 防止 finally 块再次关闭

            # 保存修改后的数据（覆盖原文件）
            np.savez_compressed(file_path, **output_data)
            print(f"已从 '{file_path}' 的 'circuits' 键中删除电路 '{key_to_delete}' 并保存更改。")
        else:
            print("操作已取消。")

    except Exception as e:
        print(f"处理文件时出错: {e}")
    finally:
        # 确保文件总是被关闭
        if original_data_handle is not None and hasattr(original_data_handle, 'close'):
             original_data_handle.close()


if __name__ == "__main__":
    # 在这里设置 NPZ 文件的路径
    NPZ_FILE_PATH = 'data/ITC22_dataset/test/circuits5.npz' 
    
    edit_npz(NPZ_FILE_PATH) 